# -*- coding:utf-8 -*-

# @Time    : 2023/5/16 02:21
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : generate_bot_dialogue.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
import pandas as pd
import asyncio
from bot.insurance_planner_gpt.planner import PlannerChat
from bot.insurance_planner_gpt.agent.user import User
from bot.insurance_planner_gpt.agent.user_suggestion import UserSuggestion
import uuid
from data_generate import utils
import traceback
from agent import llm_agent


# 生成bot对话
async def mock_bot_dialogue(path, taget_num=30):
    train_data = []
    llm_agent.change_default_llm("gpt")
    for i in range(taget_num):
        try:
            bot = PlannerChat(None)
            context = []
            instruct_dict_new = {}
            session_id = str(uuid.uuid1())
            instruct_dict_new["id"] = session_id
            instruct_dict_new["messages"] = []


            history = bot.planner.format_conversation_history(context)
            res, other_message = await bot.planner.async_reply(history, session_id, "test", is_stream=False)
            conversation_dict = {}
            conversation_dict["role"] = "assistant"
            conversation_dict["content"] = res
            instruct_dict_new["messages"].append(conversation_dict)
            context.append(conversation_dict)

            for i in range(8):
                history = bot.planner.format_conversation_history(context)
                user = User(history)
                query = await user.achat_with_proxy_gpt4(save_data=False)
                conversation_dict = {}
                conversation_dict["role"] = "user"
                conversation_dict["content"] = query.replace("<end>", "")
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

                # history = bot.planner.format_conversation_history(context)
                res, other_message = await bot.planner.async_reply(context, session_id, "test", is_stream=False)
                conversation_dict = {}
                conversation_dict["role"] = "assistant"
                conversation_dict["content"] = res
                instruct_dict_new["messages"].append(conversation_dict)
                context.append(conversation_dict)

                history = bot.planner.format_conversation_history(context)

                user_suggestion_obj = UserSuggestion(conversation_history=history)
                user_suggestion = await user_suggestion_obj.achat_auto_llm(type="gpt")
                print(user_suggestion)

                if "<end>" in query:
                    break
            train_data.append(instruct_dict_new)
            utils.jdump(train_data, path)
        except Exception as e:
            traceback.print_exc()
            # 打印堆栈
            print(e)


def add_empty_row(group):
    return pd.concat([group, pd.DataFrame([[''] * len(group.columns)], columns=group.columns)],
                     ignore_index=True)


def conversation2csv(file_path):
    datas = utils.jload(file_path)
    df = pd.DataFrame(columns=['角色', '内容'])

    for data in datas:
        messages = data['messages']
        for text in messages:
            df = df.append(
                {"角色": text['role'],
                 "内容": text['content']},
                ignore_index=True)
        df = df.reset_index(drop=True).append(
            {"角色": "",
             "内容": ""},
            ignore_index=True)
    df.to_csv("gen_conversation.csv")


if __name__ == '__main__':
    import datetime
    now_date = datetime.datetime.now().strftime("%Y-%m-%d")
    file_path = "../data_set/bot_dialogue/bot_dialogue_" + now_date +".json"  # 文件夹路径
    asyncio.run(mock_bot_dialogue(file_path))
    # conversation2csv(file_path)
